﻿
using System;
using System.Collections.Generic;
using System.Linq;
using PhoneNetTools.Dns.Records;
using PhoneNetTools.Utilities;

namespace PhoneNetTools.Dns
{
    /// <summary>
    /// Represents a DNS message containing the header data, the questions and returned resource records.
    /// </summary>
    public class Message
    {
        private const int HeaderLength = 12;
        private byte[] _rawData = new byte[HeaderLength];

        private readonly List<Question> _questions = new List<Question>();
        private readonly List<GenericResourceRecord> _answerRecords = new List<GenericResourceRecord>();
        private readonly List<GenericResourceRecord> _nameServerRecords = new List<GenericResourceRecord>();
        private readonly List<GenericResourceRecord> _additionalRecords = new List<GenericResourceRecord>();

        #region Header Fields

        /// <summary>
        /// Gets or sets the id. This identifier is copied the corresponding reply and can be used by the requester to match up replies to outstanding queries.
        /// </summary>
        public ushort Id
        {
            get
            {
                GuardHeader();
                return BigEndianBitConverter.ToUInt16(_rawData, 0);
            }
            set
            {
                GuardHeader();
                var raw = BigEndianBitConverter.GetBytes(value);
                Array.Copy(raw, 0, _rawData, 0, 2);
            }
        }

        /// <summary>
        /// Gets or sets a value indicating whether whether this message is a query (<c>false</c>), or a response (<c>true</c>).
        /// </summary>
        /// <value>
        ///   <c>true</c> if this is a response; otherwise, <c>false</c>.
        /// </value>
        public bool IsResponse
        {
            get
            {
                GuardHeader();
                return (_rawData[2] & (0x1 << 7)) > 0;
            }
            set
            {
                GuardHeader();
                _rawData[2] = value ? (byte)((0x1 << 7) | _rawData[2]) : (byte)(~(0x1 << 7) & _rawData[2]);
            }
        }

        /// <summary>
        /// Gets or sets the op code. A four bit field that specifies kind of query in this message. This value is set by the originator of a query and copied into the response.
        /// </summary>
        public OpCode OpCode
        {
            get
            {
                GuardHeader();
                return (OpCode)((_rawData[2] >> 3) & 0xf);
            }
            set
            {
                GuardHeader();
                _rawData[2] |= (byte)(((int)value & 0xf) << 3);
            }
        }

        /// <summary>
        /// Gets or sets a value indicating whether the responding name server is an authority for the domain name in the question section.
        /// This bit is valid in responses only.
        /// </summary>
        /// <value>
        /// 	<c>true</c> if the queried server is an authority for the domain name; otherwise, <c>false</c>.
        /// </value>
        public bool IsAuthoritativeAnswer
        {
            get
            {
                GuardHeader();
                return ((_rawData[2] >> 2) & 0x1) == 0x1;
            }
            set
            {
                GuardHeader();
                _rawData[2] = value ? (byte)((0x1 << 2) | _rawData[2]) : (byte)(~(0x1 << 2) & _rawData[2]);
            }
        }

        /// <summary>
        /// Gets or sets a value indicating whether this message was truncated due to length greater than that permitted on the transmission channel.
        /// </summary>
        /// <value>
        /// 	<c>true</c> if this message is truncated; otherwise, <c>false</c>.
        /// </value>
        public bool IsTruncated
        {
            get
            {
                GuardHeader();
                return ((_rawData[2] >> 1) & 0x1) == 0x1;
            }
            set
            {
                GuardHeader();
                _rawData[2] = value ? (byte)((0x1 << 1) | _rawData[2]) : (byte)(~(0x1 << 1) & _rawData[2]);
            }
        }

        /// <summary>
        /// Gets or sets a value indicating whether the name server should pursue the query recursively.
        /// Recursive query support is optional..
        /// </summary>
        /// <value>
        /// 	<c>true</c> if the name server should pursue the query recursively; otherwise, <c>false</c>.
        /// </value>
        public bool IsRecursionDesired
        {
            get
            {
                GuardHeader();
                return (_rawData[2] & 0x1) == 0x1;
            }
            set
            {
                GuardHeader();
                _rawData[2] = value ? (byte)(0x1 | _rawData[2]) : (byte)(~0x1 & _rawData[2]);
            }
        }

        /// <summary>
        /// Gets or sets a value indicating whether recursive query support is available in the name server.
        /// Set in a response by the server.
        /// </summary>
        /// <value>
        /// 	<c>true</c> if recursive query support is available in the name server; otherwise, <c>false</c>.
        /// </value>
        public bool IsRecursionAvailable
        {
            get
            {
                GuardHeader();
                return (_rawData[3] & (0x1 << 7)) > 0;
            }
            set
            {
                GuardHeader();
                _rawData[3] = value ? (byte)((0x1 << 7) | _rawData[3]) : (byte)(~(0x1 << 7) & _rawData[3]);
            }
        }

        /// <summary>
        /// Gets or sets the response code set by the name server.
        /// </summary>
        public ResponseCode ResponseCode
        {
            get
            {
                GuardHeader();
                return (ResponseCode)(_rawData[3] & 0xf);
            }
            set
            {
                GuardHeader();
                _rawData[3] |= (byte)((int)value & 0xf);
            }
        }

        private void GuardHeader()
        {
            if (_rawData == null || _rawData.Length < HeaderLength)
            {
                throw new InvalidOperationException("Not initialized or header has the wrong format/length.");
            }
        }

        #endregion

        /// <summary>
        /// Gets the collection of questions to be send to the server.
        /// </summary>
        public IList<Question> Questions
        {
            get
            {
                return _questions;
            }
        }

        /// <summary>
        /// Gets the collection of answer records returned by the server.
        /// </summary>
        public IList<GenericResourceRecord> AnswerRecords
        {
            get
            {
                return _answerRecords;
            }
        }

        /// <summary>
        /// Gets the collection of name server records returned by the server.
        /// </summary>
        public IList<GenericResourceRecord> NameServerRecords
        {
            get
            {
                return _nameServerRecords;
            }
        }

        /// <summary>
        /// Gets the collection of additional records returned by the server.
        /// </summary>
        public IList<GenericResourceRecord> AdditionalRecords
        {
            get
            {
                return _additionalRecords;
            }
        }

        /// <summary>
        /// Initializes the message contents from a raw representation.
        /// </summary>
        /// <param name="rawData">The raw data used to restore the message contents, including all header data, questions, and returned resource records.</param>
        public void InitializeFrom(byte[] rawData)
        {
            if (rawData.Length < HeaderLength)
            {
                throw new ArgumentException("Invalid raw data; header format insufficient.");
            }

            // check if mandatory bits are zero
            var isValid = ((rawData[3] >> 4) & 0x7) == 0;
            if (!isValid)
            {
                throw new ArgumentException("Z-Code is not zero. Invalid header format.");
            }

            // store (only used for the header information)
            _rawData = rawData;

            // clean all existing data
            _questions.Clear();
            _answerRecords.Clear();
            _nameServerRecords.Clear();
            _additionalRecords.Clear();

            // extract the resources
            var currentIndex = HeaderLength;

            // => questions
            var questionsCount = BigEndianBitConverter.ToUInt16(_rawData, 4);
            for (int i = 0; i < questionsCount; i++)
            {
                var question = ResourceRecordFactory.CreateQuestionFromRaw(_rawData, ref currentIndex);
                Questions.Add(question);
            }

            // => answers
            var answersCount = BigEndianBitConverter.ToUInt16(_rawData, 6);
            for (int i = 0; i < answersCount; i++)
            {
                var resourceRecord = ResourceRecordFactory.CreateFromRaw(_rawData, ref currentIndex);
                AnswerRecords.Add(resourceRecord);
            }

            // => name server count
            var nameServersCount = BigEndianBitConverter.ToUInt16(_rawData, 8);
            for (int i = 0; i < nameServersCount; i++)
            {
                var resourceRecord = ResourceRecordFactory.CreateFromRaw(_rawData, ref currentIndex);
                NameServerRecords.Add(resourceRecord);
            }

            // => additional records
            var additionalRecordsCount = BigEndianBitConverter.ToUInt16(_rawData, 10);
            for (int i = 0; i < additionalRecordsCount; i++)
            {
                var resourceRecord = ResourceRecordFactory.CreateFromRaw(_rawData, ref currentIndex);
                AdditionalRecords.Add(resourceRecord);
            }
        }

        /// <summary>
        /// Creates a raw representation of the message that can be used for data transmission.
        /// </summary>
        /// <returns>A byte array containing all the message data, including the header, questions and resource records.</returns>
        public byte[] ToRawData()
        {
            // get total length
            var length = HeaderLength
                + Questions.Sum(question => question.Length)
                + AnswerRecords.Sum(answer => answer.Length)
                + NameServerRecords.Sum(nameServerRecord => nameServerRecord.Length)
                + AdditionalRecords.Sum(additionalRecord => additionalRecord.Length);

            // set count headers
            var raw = BigEndianBitConverter.GetBytes((ushort)Questions.Count);
            Array.Copy(raw, 0, _rawData, 4, 2);
            raw = BigEndianBitConverter.GetBytes((ushort)AnswerRecords.Count);
            Array.Copy(raw, 0, _rawData, 6, 2);
            raw = BigEndianBitConverter.GetBytes((ushort)NameServerRecords.Count);
            Array.Copy(raw, 0, _rawData, 8, 2);
            raw = BigEndianBitConverter.GetBytes((ushort)AdditionalRecords.Count);
            Array.Copy(raw, 0, _rawData, 10, 2);

            // create result and copy header
            var result = new byte[length];
            Array.Copy(_rawData, 0, result, 0, HeaderLength);

            var currentIndex = HeaderLength;

            // add individual parts
            foreach (var question in Questions)
            {
                currentIndex = question.ToRawData(result, currentIndex);
            }

            foreach (var record in AnswerRecords)
            {
                currentIndex = record.ToRawData(result, currentIndex);
            }

            foreach (var record in NameServerRecords)
            {
                currentIndex = record.ToRawData(result, currentIndex);
            }

            foreach (var record in AdditionalRecords)
            {
                currentIndex = record.ToRawData(result, currentIndex);
            }

            return result;
        }
    }
}
